package org.huangrui.spark.scala.sql

import org.apache.spark.sql.expressions.Aggregator
import org.apache.spark.sql.{DataFrame, Dataset, Encoder, Encoders, SparkSession, TypedColumn, functions}

/**
 * @Author hr
 * @Create 2024-10-20 18:12 
 */
object SparkSQL03_SQL_UDAF_1 {
  def main(args: Array[String]): Unit = {
    val spark = SparkSession.builder().appName("SparkSQL03_SQL_UDAF").master("local[2]").getOrCreate()
    spark.read.json("data/user.json").createOrReplaceTempView("user")

    import spark.implicits._
    val df: DataFrame = spark.read.json("data/user.json")
    // 早期版本中，spark不能在sql中使用强类型UDAF操作
    // SQL & DSL
    // 早期的UDAF强类型聚合函数使用DSL语法操作
    val ds: Dataset[User] = df.as[User]
    // 将UDAF函数转换为查询的列对象
    val udafCol = new MyAvgUDAF().toColumn
    ds.select(udafCol).show()

    spark.stop()
  }
  case class User(id:Long,name:String, age:Long)
  case class Buff(var total: Long, var count: Long)

  /**
   * 自定义聚合函数类：计算年龄的平均值
   *      1. 继承org.apache.spark.sql.expressions.Aggregator, 定义泛型
   *         IN : 输入的数据类型 Long
   *         BUF : 缓冲区的数据类型 Buff
   *         OUT : 输出的数据类型 Long
   *         2. 重写方法(6)
   */
  class MyAvgUDAF extends Aggregator[User, Buff, Long] {
    // z & zero : 初始值或零值
    // 缓冲区的初始化
    override def zero = {
      Buff(0, 0)
    }

    // 根据输入的数据更新缓冲区的数据
    override def reduce(buff: Buff, in: User) = {
      buff.total = buff.total + in.age
      buff.count = buff.count + 1
      buff
    }

    // 合并缓冲区
    override def merge(buff1: Buff, buff2: Buff) = {
      buff1.total = buff1.total + buff2.total
      buff1.count = buff1.count + buff2.count
      buff1
    }

    //计算结果
    override def finish(buff: Buff) = {
      buff.total / buff.count
    }

    // 缓冲区的编码操作
    override def bufferEncoder = Encoders.product

    // 输出的编码操作
    override def outputEncoder = Encoders.scalaLong
  }
}
